#!/usr/bin/python -O

import collections 
import random 
import sys 

sys.path.append('src/python/modules')

import utils

if len(sys.argv) < 4:
    utils.die('usage: %s path max_rows strat_field1 strat_field2 ...' % sys.argv[0])

path = sys.argv[1]
max_rows = int(sys.argv[2])
strat_fields = sys.argv[3:]
strat_indexes = []
counts = collections.defaultdict(int)
line_num = 0

utils.warn('counting records in stratification groupings.')
for line in open(path):
    if line_num == 0:    # skip header
        fields = [f.strip() for f in line.split('\t')]
        for f in strat_fields:
            f = f.strip()
            if f in fields:
                strat_indexes.append(fields.index(f))
            else:
                utils.die('unknown field: %s (not in %s)' % (`f`, fields))
        utils.warn('stratifying on field indexes %s' % (strat_indexes))
    else:
        tokens = [f.strip() for f in line.split('\t')]
        key = tuple([tokens[i] for i in strat_indexes])
        counts[key] += 1
    line_num += 1

utils.warn('total number of records is %s' % line_num)
utils.warn('total number of stratification groups is %s' % (len(counts)))
utils.warn('number that are full is %s' % (len([v for v in counts.values() if v >= max_rows])))
utils.warn('total number of sampled records is %s' % sum([min(max_rows, n) for n in counts.values()]))

sample_indexes = {}
for (key, n) in counts.items():
    if n <= max_rows:
        sample_indexes[key] = set(xrange(n))
    else:
        sample_indexes[key] = set(random.sample(xrange(n), max_rows))
    counts[key] = 0

line_num = 0
for line in open(path):
    if line_num == 0:    # skip header
        sys.stdout.write(line)
    else:
        tokens = [f.strip() for f in line.split('\t')]
        key = tuple([tokens[i] for i in strat_indexes])
        if counts[key] in sample_indexes[key]:
            sys.stdout.write(line)
        counts[key] += 1
    line_num += 1
